import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory

# --- 环境准备 ---
print("正在从 env.example 加载环境变量...")
load_dotenv(dotenv_path="env.example")
api_key = os.getenv("OPENAI_API_KEY")
model_name = os.getenv("OPENAI_MODEL", "qwen-plus-latest")

if not api_key:
    print("错误：请先在 env.example 文件中设置你的 OPENAI_API_KEY。")
else:
    print("环境变量加载成功！")
    # --- 核心代码开始 (LCEL 现代方法) ---

    # 1. 初始化模型
    llm = ChatOpenAI(model=model_name)
    print(f"模型已初始化 (使用 {model_name})。")

    # 2. 创建一个无状态的核心链 (prompt | llm | parser)
    prompt = ChatPromptTemplate.from_messages([
        ("system", "你是一个风趣幽默的AI助手。"),
        MessagesPlaceholder(variable_name="chat_history"), # 占位符，用于插入历史记录
        ("user", "{input}"),
    ])
    output_parser = StrOutputParser()
    chain = prompt | llm | output_parser
    print("核心链已创建。")

    # 3. 创建一个用于存储历史记录的 "数据库"
    # 在真实应用中，这里应该是 Redis, SQL, etc.
    # 在这个例子里，我们用一个简单的字典来模拟。
    store = {}

    # 4. 创建一个函数，用于根据 session_id 获取对应的历史记录
    def get_session_history(session_id: str) -> BaseChatMessageHistory:
        if session_id not in store:
            store[session_id] = ChatMessageHistory() # 如果是新会话，创建一个新的历史记录对象
        return store[session_id]

    # 5. 用 RunnableWithMessageHistory "包装" 核心链
    # 这就是魔法发生的地方！
    with_message_history = RunnableWithMessageHistory(
        chain,
        get_session_history,
        input_messages_key="input",        # 告诉它用户的输入在哪个字段
        history_messages_key="chat_history", # 告诉它历史记录应该插入到提示的哪个占位符
    )
    print("带记忆功能的链已创建！")

    # 6. 开始对话！
    # 我们需要为这次对话指定一个唯一的 ID
    session_id = "my_first_chat"

    print("\n--- 开始与 AI 对话 (LCEL 方式) ---")
    
    # 第一次对话
    print("\n[你]: 你好，我叫小明。")
    response_1 = with_message_history.invoke(
        {"input": "你好，我叫小明。"},
        config={"configurable": {"session_id": session_id}} # 必须传入 config
    )
    print(f"[AI]: {response_1}")

    # 第二次对话
    print("\n[你]: 我叫什么名字？")
    response_2 = with_message_history.invoke(
        {"input": "我叫什么名字？"},
        config={"configurable": {"session_id": session_id}} # 传入相同的 session_id
    )
    print(f"[AI]: {response_2}")
    
    print("\n--- 对话结束 ---")

    # 你可以手动查看 store 中存储的内容
    print("\n--- 查看 Memory Store 中的历史记录 ---")
    print(f"会话 '{session_id}' 的历史: {store[session_id].messages}")
    print("--------------------") 